import json
import os

import pandas as pd
from dotenv import load_dotenv
from langchain_core.tools import tool
from pydantic import BaseModel, Field
import pymysql

description = """
"当用户需要进行数据库查询工作时，请调用该函数，该函数用于在指定mysql服务器运行一段SQL代码，完成数据查询的相关工作
并且当前函数是使用pymysql连接Mysql数据库.
本函数只负责运行SQL代码并进行数据查询，若要进行数据提取则使用另一个extract_data函数
"""
class SQLQuerySchema(BaseModel):
    sql_query: str = Field(description=description)


@tool(args_schema=SQLQuerySchema)
def sql_inter(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数，该函数用于在指定mysql服务器运行一段SQL代码，完成数据查询的相关工作
    并且当前函数是使用pymysql连接Mysql数据库.
    本函数只负责运行SQL代码并进行数据查询，若要进行数据提取则使用另一个extract_data函数
    :param sql_query: 要执行的SQL查询语句
    :return: sql_query在mysql中运行的结果
    """
    load_dotenv(override=True)
    host = os.getenv("MYSQL_HOST")
    user = os.getenv("MYSQL_USER")
    mysql_pw = os.getenv("MYSQL_PW")
    db = os.getenv("MYSQL_DB")
    port = int(os.getenv("MYSQL_PORT"))

    # 连接到MySQL数据库
    connection = pymysql.connect(
        host=host,
        user=user,
        password=mysql_pw,
        database=db,
        port=port,
        charset='utf8'
    )

    try:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            result = cursor.fetchall()
            return json.dumps(result,default=str)

    finally:
        connection.close()


class ExtractQuerySchema(BaseModel):
    sql_query: str = Field(description="用于从Mysql提取数据的SQL查询语句")
    df_name: str = Field(description="用于保存结果的pandas变量名词(字符串形式)")


@tool(args_schema=ExtractQuerySchema)
def extract_data(sql_query: str, df_name: str) -> str:
    """
    用于在Mysql数据库中提取一张表到当前python环境中，注意，本函数只负责数据表的提取，并不负责数据查询。
    若需要在Mysql中进行数据查询，请使用sql_inter函数
    同时需要注意，编写外部函数的参数消息时，必须要满足json格式的字符串""
    :param sql_query: 字符串形式的SQL查询语句，用于提取mysql中的某张表
    :param df_name : 将Mysql数据库中提取的表格进行本地保存时的变量名，以字符串形式表示
    :return: 表格读取和保存结果
    """
    print("正在调用extract_data工具sql查询...")
    load_dotenv(override=True)
    host = os.getenv("MYSQL_HOST")
    user = os.getenv("MYSQL_USER")
    mysql_pw = os.getenv("MYSQL_PW")
    db = os.getenv("MYSQL_DB")
    port = int(os.getenv("MYSQL_PORT"))

    # 连接到MySQL数据库
    connection = pymysql.connect(
        host=host,
        user=user,
        password=mysql_pw,
        database=db,
        port=port,
        charset='utf8'
    )

    try:
        # 执行SQL并保存为全局变量
        df = pd.read_sql(sql_query, connection)
        globals()[df_name] = df
        print(f"数据已保存到变量 {df_name} 中")
        return f"成功创建pandas对象{df_name} ,包含从Mysql提取的数据"
    except Exception as e:
        return f"提取数据时出错: {e}"
    finally:
        connection.close()
